package utils

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

// SDProxy SD API 代理工具
type SDProxy struct {
	BaseURL    string        // SD WebUI API 基础地址
	Timeout    time.Duration // 请求超时时间
	client     *http.Client
}

// NewSDProxy 创建新的 SD API 代理实例
func NewSDProxy(baseURL string, timeout time.Duration) *SDProxy {
	if timeout == 0 {
		timeout = 5 * time.Minute // 默认超时时间 5 分钟
	}
	return &SDProxy{
		BaseURL: baseURL,
		Timeout: timeout,
		client: &http.Client{
			Timeout: timeout,
		},
	}
}

// ProxyRequest 通用的请求转发方法
func (p *SDProxy) ProxyRequest(ctx context.Context, endpoint string, reqBody interface{}, respBody interface{}) error {
	// 1. 构建完整 URL
	url := fmt.Sprintf("%s%s", p.BaseURL, endpoint)

	// 2. 将请求体转换为 JSON
	var bodyReader io.Reader
	if reqBody != nil {
		jsonData, err := json.Marshal(reqBody)
		if err != nil {
			return fmt.Errorf("marshal request body failed: %w", err)
		}
		bodyReader = bytes.NewBuffer(jsonData)
	}

	// 3. 创建请求
	req, err := http.NewRequestWithContext(ctx, "POST", url, bodyReader)
	if err != nil {
		return fmt.Errorf("create request failed: %w", err)
	}

	// 4. 设置请求头
	req.Header.Set("Content-Type", "application/json")

	// 5. 发送请求
	resp, err := p.client.Do(req)
	if err != nil {
		return fmt.Errorf("send request failed: %w", err)
	}
	defer resp.Body.Close()

	// 6. 检查响应状态
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("SD API error: status=%d, body=%s", resp.StatusCode, string(body))
	}

	// 7. 解析响应体
	if respBody != nil {
		if err := json.NewDecoder(resp.Body).Decode(respBody); err != nil {
			return fmt.Errorf("decode response failed: %w", err)
		}
	}

	return nil
}

// Txt2Img 文生图接口
func (p *SDProxy) Txt2Img(ctx context.Context, req interface{}) (*Txt2ImgResponse, error) {
	resp := &Txt2ImgResponse{}
	err := p.ProxyRequest(ctx, "/sdapi/v1/txt2img", req, resp)
	return resp, err
}

// Img2Img 图生图接口
func (p *SDProxy) Img2Img(ctx context.Context, req interface{}) (*Img2ImgResponse, error) {
	resp := &Img2ImgResponse{}
	err := p.ProxyRequest(ctx, "/sdapi/v1/img2img", req, resp)
	return resp, err
}

// GetProgress 获取进度接口
func (p *SDProxy) GetProgress(ctx context.Context) (*ProgressResponse, error) {
	resp := &ProgressResponse{}
	err := p.ProxyRequest(ctx, "/sdapi/v1/progress", nil, resp)
	return resp, err
}

// 响应结构体定义
type Txt2ImgResponse struct {
	Images     []string               `json:"images"`
	Parameters map[string]interface{} `json:"parameters"`
	Info       string                 `json:"info"`
}

type Img2ImgResponse struct {
	Images     []string               `json:"images"`
	Parameters map[string]interface{} `json:"parameters"`
	Info       string                 `json:"info"`
}

type ProgressResponse struct {
	Progress     float64 `json:"progress"`
	EtaRelative  float64 `json:"eta_relative"`
	State        State   `json:"state"`
	CurrentImage string  `json:"current_image"`
}

type State struct {
	Skipped       bool    `json:"skipped"`
	Interrupted   bool    `json:"interrupted"`
	Job           string  `json:"job"`
	JobCount      int     `json:"job_count"`
	JobTimestamp  string  `json:"job_timestamp"`
	JobNo         int     `json:"job_no"`
	Sampling      string  `json:"sampling_step"`
	SamplingSteps int     `json:"sampling_steps"`
}
